"""
联邦学习框架实现
基于Flower框架实现联邦学习功能
"""

import asyncio
from typing import List, Dict, Any, Callable, Optional, Tuple
from flwr.common import Metrics, Parameters, Scalar
from flwr.server.strategy import FedAvg
from flwr.server.client_manager import ClientManager, SimpleClientManager
import flwr as fl
import numpy as np
from src.research_core.model_manager import ModelManager, ModelType


class CustomStrategy(FedAvg):
    """自定义联邦学习策略"""
    
    def aggregate_fit(
        self,
        server_round: int,
        results,
        failures,
    ) -> Tuple[Optional[Parameters], Dict[str, Scalar]]:
        """聚合来自客户端的权重"""
        aggregated_parameters, aggregated_metrics = super().aggregate_fit(server_round, results, failures)
        
        if aggregated_parameters is not None:
            print(f"第 {server_round} 轮训练完成，聚合权重已更新")
        
        return aggregated_parameters, aggregated_metrics


class LangChainFederatedServer:
    """LangChain联邦学习服务器"""
    
    def __init__(self, model_type: ModelType, min_fit_clients: int = 2, min_available_clients: int = 2):
        self.model_type = model_type
        self.min_fit_clients = min_fit_clients
        self.min_available_clients = min_available_clients
        self.strategy = CustomStrategy(
            fraction_fit=1.0,  # 所有客户端参与每轮训练
            fraction_evaluate=1.0,  # 所有客户端参与每轮评估
            min_fit_clients=min_fit_clients,
            min_available_clients=min_available_clients,
            min_evaluate_clients=min_fit_clients,
        )
        self.client_manager = SimpleClientManager()
        
    def start_server(self, server_address: str = "0.0.0.0:8080"):
        """启动联邦学习服务器"""
        fl.server.start_server(
            server_address=server_address,
            config=fl.server.ServerConfig(num_rounds=5),
            strategy=self.strategy,
            client_manager=self.client_manager,
        )

    def get_model_manager(self) -> ModelManager:
        """获取模型管理器"""
        return ModelManager()